# coding=utf8

import binascii

__author__ = "zouxiaoliang"


class RC4Key:
    """
    RC4流加密秘钥
    """
    def __init__(self):
        self.x = 0
        self.y = 0
        self.box = []
        pass

    def set_key(self, key):
        self.x = 0
        self.box = range(256)
        for i in range(256):
            self.x = (self.x + self.box[i] + ord(key[i % len(key)])) % 256
            self.box[i], self.box[self.x] = self.box[self.x], self.box[i]

        self.x = self.y = 0


def _crypt(data, rc4_key):
    """
    rc4 编码函数
    :param data: 待编码内容
    :param rc4_key: 交换秘钥
    :return: rc4编码后的字符串
    """
    if not isinstance(rc4_key, RC4Key):
        return

    out = []
    for char in data:
        rc4_key.x = (rc4_key.x + 1) % 256
        rc4_key.y = (rc4_key.y + rc4_key.box[rc4_key.x]) % 256
        rc4_key.box[rc4_key.x], rc4_key.box[rc4_key.y] = rc4_key.box[rc4_key.y], rc4_key.box[rc4_key.x]
        out.append(chr(ord(char) ^ rc4_key.box[(rc4_key.box[rc4_key.x] + rc4_key.box[rc4_key.y]) % 256]))

    return ''.join(out)


def encode(data, key, encode_func=binascii.b2a_hex):
    """
    rc4加密
    :param data: 待加密内容
    :param key: 秘钥
    :param encode_func: 编码函数
    :return: 加密后字符串
    """
    data = _crypt(data, key)
    if encode_func:
        data = encode_func(data)
    return data


def decode(data, key, decode_func=binascii.a2b_hex):
    """
    rc4解码
    :param data: 待解密字符串
    :param key: 秘钥
    :param decode_func: 解码函数
    :return: 原始消息内容
    """
    if decode_func:
        data = decode_func(data)
    return _crypt(data, key)


def _main():
    # 需要加密的数据
    data = 'hello python'
    # 密钥
    send_key = RC4Key()
    send_key.set_key("123456")

    recv_key = RC4Key()
    recv_key.set_key("123456")
    # 加码
    encoded_data1 = encode(data=data, key=send_key)
    print encoded_data1
    # 解码
    # decoded_data = decode(data=encoded_data, key=recv_key)
    # print decoded_data

    data = 'hello python2'
    # 加码
    encoded_data2 = encode(data=data, key=send_key)
    print encoded_data2
    # 解码
    decoded_data = decode(data=encoded_data1 + encoded_data2, key=recv_key)
    print decoded_data

if __name__ == '__main__':
    _main()
